import backtrader as bt
from loguru import logger
import numpy as np
from datetime import datetime
import json
from pathlib import Path

class SKDJ(bt.Indicator):
    """
    SKDJ指标 (随机指标) - 按照广发证券公式实现
    参数:
        - period: N值，计算最高最低价的周期，默认9
        - slowk_period: M值，RSV的EMA周期，默认3
        - slowd_period: M值，K值的MA周期，默认3
    
    计算公式:
        LOWV = LLV(LOW, N)
        HIGHV = HHV(HIGH, N)
        RSV = EMA((CLOSE-LOWV)/(HIGHV-LOWV)*100, M)
        K = EMA(RSV, M)
        D = MA(K, M)
        J = 3*K - 2*D
    """
    lines = ('k', 'd', 'j')
    params = (
        ('period', 9),
        ('slowk_period', 3),
        ('slowd_period', 3),
    )
    
    def __init__(self):
        # 计算最高价和最低价
        self.highest = bt.indicators.Highest(self.data.high, period=self.params.period)
        self.lowest = bt.indicators.Lowest(self.data.low, period=self.params.period)
        
        # 计算RSV原始值
        rsv_raw = (self.data.close - self.lowest) / (self.highest - self.lowest) * 100
        
        # 计算RSV (EMA(RSV原始值, M))
        self.rsv = bt.indicators.EMA(rsv_raw, period=self.params.slowk_period)
        
        # 计算K值 (RSV的指数移动平均EMA)
        self.k = bt.indicators.EMA(self.rsv, period=self.params.slowk_period)
        
        # 计算D值 (K值的简单移动平均MA)
        self.d = bt.indicators.SimpleMovingAverage(self.k, period=self.params.slowd_period)
        
        # 计算J值
        self.j = 3 * self.k - 2 * self.d
        
        # 将计算结果赋值给lines
        self.lines.k = self.k
        self.lines.d = self.d
        self.lines.j = self.j

class MTM(bt.Indicator):
    """
    MTM指标 (动量指标)
    参数:
        - period: 计算周期，默认10
    """
    lines = ('mtm', 'mtmma')
    params = (
        ('period', 10),
        ('ma_period', 10),
    )
    
    def __init__(self):
        # 计算MTM值
        self.mtm = self.data.close - self.data.close(-self.params.period)
        
        # 计算MTM的移动平均
        self.mtmma = bt.indicators.SimpleMovingAverage(self.mtm, period=self.params.ma_period)
        
        # 将计算结果赋值给lines
        self.lines.mtm = self.mtm
        self.lines.mtmma = self.mtmma

class LWR(bt.Indicator):
    """
    LWR指标 (威廉指标变种)
    参数:
        - period: 计算周期，默认9
        - slowk_period: K值的平滑周期，默认3
        - slowd_period: D值的平滑周期，默认3
    
    计算公式:
        LOWV = LLV(LOW, N)
        HIGHV = HHV(HIGH, N)
        RSV = (HIGHV - CLOSE) / (HIGHV - LOWV) * 100
        K = SMA(RSV, M)
        D = SMA(K, M)
    """
    lines = ('k', 'd')
    params = (
        ('period', 9),
        ('slowk_period', 3),
        ('slowd_period', 3),
    )
    
    def __init__(self):
        # 计算最高价和最低价
        self.highest = bt.indicators.Highest(self.data.high, period=self.params.period)
        self.lowest = bt.indicators.Lowest(self.data.low, period=self.params.period)
        
        # 计算RSV原始值 (注意LWR的RSV计算与SKDJ相反)
        rsv_raw = (self.highest - self.data.close) / (self.highest - self.lowest) * 100
        
        # 计算K值 (RSV的平滑移动平均SMA)
        self.k = bt.indicators.SmoothedMovingAverage(rsv_raw, period=self.params.slowk_period)
        
        # 计算D值 (K值的平滑移动平均SMA)
        self.d = bt.indicators.SmoothedMovingAverage(self.k, period=self.params.slowd_period)
        
        # 将计算结果赋值给lines
        self.lines.k = self.k
        self.lines.d = self.d

def a_cross_b_from_below(a, b):
    return a[-1] < b[-1] and a[0] > b[0]
def a_cross_b_from_above(a, b):
    return a[-1] > b[-1] and a[0] < b[0]

class BossWongStrategy(bt.Strategy):
    """
    王总策略 - 打印SKDJ和MTM指标线
    """
    params = (
        ('target_trade_date', "2099-03-01"),
        ('print_interval', 5),  # 每5个交易日打印一次
        ('skdj_period', 9),
        ('skdj_slowk_period', 3),
        ('skdj_slowd_period', 3),
        ('mtm_period', 12),
        ('mtm_ma_period', 6),
        ('lwr_period', 9),
        ('lwr_slowk_period', 3),
        ('lwr_slowd_period', 3),
    )
    
    def __init__(self):
        if type(self.params.target_trade_date) == str:
            self.params.target_trade_date = datetime.strptime(self.params.target_trade_date, '%Y-%m-%d').date()
        self.stocks = self.datas
        self.print_counter = 0
        
        # 为每个股票数据添加SKDJ、MTM和LWR指标
        for data in self.stocks:
            # 添加LWR指标
            data.lwr = LWR(
                data,
                period=self.params.lwr_period,
                slowk_period=self.params.lwr_slowk_period,
                slowd_period=self.params.lwr_slowd_period
            )

            data.skdj = SKDJ(
                data, 
                period=self.params.skdj_period,
                slowk_period=self.params.skdj_slowk_period,
                slowd_period=self.params.skdj_slowd_period
            )
            data.mtm = MTM(
                data,
                period=self.params.mtm_period,
                ma_period=self.params.mtm_ma_period
            )
        
        logger.info(f"策略初始化完成，共加载 {len(self.stocks)} 只股票")
        logger.info(f"SKDJ参数: period={self.params.skdj_period}, slowk={self.params.skdj_slowk_period}, slowd={self.params.skdj_slowd_period}")
        logger.info(f"MTM参数: period={self.params.mtm_period}, ma_period={self.params.mtm_ma_period}")
        logger.info(f"LWR参数: period={self.params.lwr_period}, slowk={self.params.lwr_slowk_period}, slowd={self.params.lwr_slowd_period}")
    
    
    def stop(self):
        """策略结束时执行"""
        logger.info("策略执行完成")
        
        # 打印最终统计
        logger.debug(f"\n{'='*60}")
        logger.debug("最终指标统计")
        logger.debug(f"{'='*60}")
        
        count = 0
        count_list = []
        for data in self.stocks:
            if len(data) > 0:
                try:
                    final_price = data.close[0]
                    final_skdj_k = data.skdj.lines.k[0]
                    final_skdj_d = data.skdj.lines.d[0]
                    final_mtm = data.mtm.lines.mtm[0]
                    final_mtm_ma = data.mtm.lines.mtmma[0]
                    final_lwr_k = data.lwr.lines.k[0]
                    final_lwr_d = data.lwr.lines.d[0]
                    
                    logger.debug(f"{data._name}: 价格={final_price:.2f}, SKDJ(K/D)={final_skdj_k:.2f}/{final_skdj_d:.2f}, MTM={final_mtm:.2f}, MTM_MA={final_mtm_ma:.2f}, LWR(K/D)={final_lwr_k:.2f}/{final_lwr_d:.2f}")
                    logger.debug(f"lwr_k={final_lwr_k}, lwr_d={final_lwr_d}")

                    if a_cross_b_from_below(data.skdj.lines.k, data.skdj.lines.d) and a_cross_b_from_below(data.mtm.lines.mtm, data.mtm.lines.mtmma) and a_cross_b_from_above(data.lwr.lines.k, data.lwr.lines.d):
                        count += 1
                        count_list.append(data._name)

                except IndexError:
                    logger.warning(f"股票 {data._name} 最终数据获取失败")
        
        logger.info(f"{'='*60}")
        logger.info(f"符合条件股票数量: {count}")
        logger.info(f"符合条件股票列表: {count_list}")
        
        # 把count和count_list写入文件，写到result2.json，路径是{project_root}/data/result2.json
        project_root = Path(__file__).parent.parent.parent
        result_file = project_root / 'data' / 'result2.json'
        result_file.parent.mkdir(exist_ok=True)  # 确保data目录存在
        
        # 创建结果数据
        result_data = {
            "target_trade_date": str(self.params.target_trade_date),
            "strategy": "BOSS_WONG：skdj的k线下往上穿过d线，mtm线下往上穿过mtm的ma线，lwr的k线上往下穿过d线",
            "count": count,
            "count_list": count_list
        }
        
        with open(result_file, 'w', encoding='utf-8') as f:
            json.dump(result_data, f, ensure_ascii=False, indent=2)
        
        logger.info(f"结果已写入文件: {result_file}")

# 使用示例
if __name__ == "__main__":
    # 创建Cerebro引擎
    cerebro = bt.Cerebro()
    
    # 添加策略
    cerebro.addstrategy(BossWongStrategy)
    
    # 这里需要添加数据源
    # data = bt.feeds.YourDataFeed(...)
    # cerebro.adddata(data)
    
    # 运行回测
    # cerebro.run()
    
    print("SKDJ和MTM策略已创建，请在主程序中添加数据源并运行")
